// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/common/process/thread_debug_info.h>
#include <pollux/testing/gtest_utils.h>
#include <pollux/testing/exec/util/assert_query_builder.h>
#include <pollux/testing/exec/util/operator_test_base.h>
#include <pollux/plan/plan_builder.h>
#include <pollux/functions/udf.h>

#include <melon/unit.h>
#include <melon/init/init.h>
#include <gtest/gtest.h>

using namespace kumo::pollux;
using namespace kumo::pollux::exec::test;

// Only build if address sanitizer is not enabled since this test initiated a
// segfault which should be captured by the signal handler but gets caught by
// the address sanitizer handler instead which print a different error log.

// For clang
#ifdef __has_feature
#define IS_BUILDING_WITH_ASAN() __has_feature(address_sanitizer)
#define IS_BUILDING_WITH_TSAN() __has_feature(thread_sanitizer)
#else
// For GCC
#if defined(__SANITIZE_ADDRESS__) && __SANITIZE_ADDRESS__
#define IS_BUILDING_WITH_ASAN() 1
#else
#define IS_BUILDING_WITH_ASAN() 0
#endif
#if defined(__SANITIZE_THREAD__) && __SANITIZE_THREAD__
#define IS_BUILDING_WITH_TSAN() 1
#else
#define IS_BUILDING_WITH_TSAN() 0
#endif
#endif

#if IS_BUILDING_WITH_TSAN() || IS_BUILDING_WITH_ASAN()
#define IS_BUILDING_WITH_SAN
#endif

// Ensure the test class name has "DeathTest" as a prefix to ensure this runs in
// a single thread since the ASSERT_DEATH forks the process.
class ThreadDebugInfoDeathTest : public OperatorTestBase {};

namespace {

template <typename T>
struct InduceSegFaultFunction {
  template <typename TResult, typename TInput>
  void call(TResult& out, const TInput& in) {
    int* nullpointer = nullptr;
    *nullpointer = 6;
  }
};
} // namespace

DEBUG_ONLY_TEST_F(ThreadDebugInfoDeathTest, withinSeperateDriverThread) {
  auto vector = make_row_vector({make_flat_vector<int64_t>({1, 2, 3, 4, 5, 6})});
  register_function<InduceSegFaultFunction, int64_t, int64_t>({"segFault"});
  auto op = PlanBuilder().values({vector}).project({"segFault(c0)"}).planNode();
#ifndef IS_BUILDING_WITH_SAN
  ASSERT_DEATH(
      (assertQuery(op, vector)),
      ".*Fatal signal handler. Query Id= TaskCursorQuery_0 Task Id= test_cursor 1.*");
#endif
}

DEBUG_ONLY_TEST_F(ThreadDebugInfoDeathTest, withinQueryCompilation) {
  auto vector = make_row_vector({make_flat_vector<int64_t>({1, 2, 3, 4, 5, 6})});
  register_function<InduceSegFaultFunction, int64_t, int64_t>({"segFault"});
  // Call expression with a constant to trigger the constant folding during
  // compilation.
  auto op = PlanBuilder().values({vector}).project({"segFault(1)"}).planNode();

#ifndef IS_BUILDING_WITH_SAN
  ASSERT_DEATH(
      (assertQuery(op, vector)),
      ".*Fatal signal handler. Query Id= TaskCursorQuery_0 Task Id= test_cursor 1.*");
#endif
}

DEBUG_ONLY_TEST_F(ThreadDebugInfoDeathTest, withinTheCallingThread) {
  auto vector = make_row_vector({make_flat_vector<int64_t>({1, 2, 3, 4, 5, 6})});
  register_function<InduceSegFaultFunction, int64_t, int64_t>({"segFault"});
  auto plan =
      PlanBuilder().values({vector}).project({"segFault(c0)"}).planFragment();

  auto queryCtx = core::QueryCtx::create(
      executor_.get(),
      core::QueryConfig({}),
      std::unordered_map<std::string, std::shared_ptr<config::ConfigBase>>{},
      cache::AsyncDataCache::getInstance(),
      nullptr,
      nullptr,
      "TaskCursorQuery_0");
  auto task = exec::Task::create(
      "single.execution.task.0",
      std::move(plan),
      0,
      queryCtx,
      exec::Task::ExecutionMode::kSerial);

#ifndef IS_BUILDING_WITH_SAN
  ASSERT_DEATH(
      (task->next()),
      ".*Fatal signal handler. Query Id= TaskCursorQuery_0 Task Id= single.execution.task.0.*");
#endif
}

DEBUG_ONLY_TEST_F(ThreadDebugInfoDeathTest, noThreadContextSet) {
  int* nullpointer = nullptr;
#ifndef IS_BUILDING_WITH_SAN
  ASSERT_DEATH((*nullpointer = 6), ".*ThreadDebugInfo object not found.*");
#endif
  melon::compiler_must_not_elide(nullpointer);
}
